import pandas as pd
import statsmodels.formula.api as smf
from sklearn.model_selection import train_test_split
from sqlalchemy import create_engine
from sklearn.metrics import r2_score, mean_absolute_error
from matplotlib import pyplot as plt
import numpy as np
import seaborn as sns
import pymysql

db_config = {
    "host": "127.0.0.1",
    "user": "root",
    "password": "root",
    "database": "yqt",
    "port": 3306,
    "charset": "utf8"  # 字符集设置
}

engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
)
conn = pymysql.connect(**db_config)
chunk_size = 10000

df = pd.read_sql_query(
    """
    SELECT d.*, m.buy_lg_vol, m.sell_lg_vol, m.buy_alg_vol, m.sell_alg_vol, m.net_mf_vol, i.vol as i_vol, i.closes as i_closes
    FROM date_1 d JOIN moneyflows m ON d.ts_code = m.ts_code and d.trade_date = m.trade_date
    LEFT JOIN index_daily i ON i.trade_date = d.trade_date and i.ts_code = '999908.SZ'
    """
)
df1['zd_closes'] = round((df1['i_closes'] - df1['i_closes'].shift(1)) / df1['i_closes'].shift(1), 2)
df1['zs_vol'] = round((df1['i_vol'] - df1['i_vol'].shift(1)) / df1['i_vol'].shift(1), 2)

# 处理缺失数据
df1 = df1.dropna(subset=['zd_closes', 'zs_closes', 'zs_vol'])
print(df1.head())

ex = ['id', 'ts_code', 'trade_date', 'the_date', 'opens', 'high', 'low', 'closes', 'pre_closes', 'changes', 'pct_chg', 'amount', 'zd_closes', 'zs_vol', 'i_vol']
number = df1.select_dtypes(include=['number']).columns.to_list()
nealist = [col for col in number if col not in ex]

# 回归分析
formula = 'zd_closes ~'+ '+'.join(nealist)
res = smf.ols(formula, data=df1).fit()
print(res.summary())

# 可视化
plt.figure(figsize=(10, 6))
plt.scatter(res.fittedvalues, df1['zd_closes'])
plt.show()

features = df1[['i_vol', 'buy_lg_vol','sell_lg_vol', 'buy_alg_vol','sell_alg_vol', 'net_mf_vol', 'i_closes', 'zs_closes']]
correlation_matrix = features.corr()

plt.figure(figsize=(10, 6))
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', cbar=True, fmt='.2f', linewidths=.5)
plt.show()